/******************************************************************************
	This file is a part of the Snapshot Bioinformatics Project

	Copyright (C) 2012-2013 Jarom Schow

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>.
******************************************************************************/

#include "GenBankParser.h"

#include <QtCore/QIODevice>
#include <QtCore/QFile>

#include "Core.h"
#include "FileUtil.h"
#include "GzFile.h"
#include "Strings.h"
#include "SeqFeature.h"
#include "SeqReference.h"
#include "SeqRange.h"
#include "Progress.h"

using namespace Bio;

const int	c_bufferSize = 128;

const int	c_headerOffset = 12;
const int	c_featureOffset = 21;
const int	c_originOffset = 10;

inline bool	isSpace( char ch )
{
	return (ch == ' ' || ch == '\t' || ch == '\r' || ch == '\n');
}

/**
**/
GenBankParser::GenBankParser() : m_parseMode( None ), m_featureMode( Invalid ), 
	m_curRecord( NULL ), m_curFeature( NULL ), m_curReference( NULL )
{
}

/**
**/
GenBankParser::~GenBankParser()
{
	delete m_curRecord;
	delete m_curFeature;
	delete m_curReference;
}

/**
**/
bool GenBankParser::parse( QIODevice &dev )
{
	if( !dev.isOpen() )
		return false;

	const QIODevice	*sourceDev = &dev;
	QString			filename = "GenBank data";

	if( const GzFile *gzFile = qobject_cast<const GzFile*>( sourceDev ) )
		sourceDev = gzFile->getSourceDevice();

	if( const QFile *file = qobject_cast<const QFile*>( sourceDev ) )
		filename = FileUtil::getFilename( file->fileName() );

	char		line[c_bufferSize];
	Progress	progress( QString( "Parsing %1..." ).arg( filename ), int( sourceDev->size() ) );

	m_parseMode = None;
	m_featureMode = Invalid;
	m_lineNo = 0;

	while( !dev.atEnd() ){
		dev.readLine( line, c_bufferSize );
		m_lineNo++;
		progress.update( int( sourceDev->pos() ) );

		if( progress.isCanceled() )
			return false;

		if( m_parseMode == None ){
			if( !strncmp( "LOCUS", line, 5 ) ){
				startRecord();
				startSection( Locus );
				parseLocus( line );
			}
		}
		else {
			if( line[0] == ' ' ){
				switch( m_parseMode ){
				case Locus:
					parseLocus( line );
					break;
				case Version:
					parseVersion( line );
					break;
				case Source:
					parseSource( line );
					break;
				case Organism:
					parseOrganism( line );
					break;
				case Reference:
				case Authors:
				case Title:
				case Journal:
				case Pubmed:
					parseReference( line );
					break;
				case Features:
					parseFeatures( line );
					break;
				case Origin:
					addText( line, c_originOffset );
					break;
				default:
					addText( line, c_headerOffset );
					break;
				}
			}
			else if( line[0] == '/' ) {
				if( line[1] == '/' )
					endRecord();
			}
			else if( !strncmp( "DEFINITION", line, 10 ) ) {
				startSection( Definition );
				addText( line, c_headerOffset );
			}
			else if( !strncmp( "ACCESSION", line, 9 ) ) {
				startSection( Accession );
				addText( line, c_headerOffset );
			}
			else if( !strncmp( "VERSION", line, 7 ) ) {
				startSection( Version );
				parseVersion( line );
			}
			else if( !strncmp( "DBLINK", line, 6 ) ) {
				startSection( DBLink );
				addText( line, c_headerOffset );
			}
			else if( !strncmp( "KEYWORDS", line, 8 ) ) {
				startSection( Keywords );
				addText( line, c_headerOffset );
			}
			else if( !strncmp( "SOURCE", line, 6 ) ) {
				startSection( Source );
				parseSource( line );
			}
			else if( !strncmp( "REFERENCE", line, 9 ) ) {
				startSection( Reference );
				startReference();
			}
			else if( !strncmp( "COMMENT", line, 7 ) ) {
				startSection( Comment );
				addText( line, c_headerOffset );
			}
			else if( !strncmp( "FEATURES", line, 8 ) ) {
				if( !strncmp( "Location/Qualifiers", line + 21, 19 ) )
					startSection( Features );
				else
					startSection( None ); // skip it if the features are not Location/Qualifiers
			}
			else if( !strncmp( "ORIGIN", line, 6 ) ) {
				startSection( Origin );
			}
		}
	}

	progress.finish();

	return true;
}

/**
**/
int GenBankParser::getNumRecords() const
{
	return m_seqRecords.count();
}

/**
**/
SeqRecord& GenBankParser::getRecord( int which )
{
	if( which < 0 || which >= m_seqRecords.count() ){
		static SeqRecord	empty;

		return empty;
	}

	return m_seqRecords[which];
}

/**
**/
SeqRecord GenBankParser::getRecord( int which ) const
{
	if( which < 0 || which >= m_seqRecords.count() )
		return SeqRecord();

	return m_seqRecords[which];
}

/**
**/
QMutableListIterator<SeqRecord> GenBankParser::getMutableRecordIterator()
{
	return QMutableListIterator<SeqRecord>( m_seqRecords );
}

/**
**/
QListIterator<SeqRecord> GenBankParser::getRecordIterator() const
{
	return QListIterator<SeqRecord>( m_seqRecords );
}

/**
**/
void GenBankParser::lineError( const QString &msg )
{
	error( QString( "Error on line %1: " ).arg( m_lineNo ) + msg );
}

/**
**/
void GenBankParser::lineWarning( const QString &msg )
{
	warning( QString( "Warning on line %1: " ).arg( m_lineNo ) + msg );
}

/**
**/
bool GenBankParser::startRecord()
{
	if( m_curRecord )
		return false;

	m_curRecord = new SeqRecord;

	return true;
}

/**
**/
bool GenBankParser::endRecord()
{
	if( !m_curRecord )
		return false;

	startSection( None ); // flush any data that has been read for the current record

	if( m_curReference )
		endReference();

	m_seqRecords.append( *m_curRecord );
	delete m_curRecord;
	m_curRecord = NULL;

	return true;
}

/**
**/
bool GenBankParser::startSection( ParseMode mode )
{
	if( !m_curRecord )
		return false;

	if( !m_text.isEmpty() ) {
		switch( m_parseMode ){
		case Definition:
			m_curRecord->setDescription( m_text );
			break;
		case Accession: {
			int	idx = m_text.indexOf( ' ' );

			if( idx > 0 ) // only pay attention to the first accession - todo: record the others
				m_text = m_text.left( idx );

			m_curRecord->setId( m_text );
			break; }
		case DBLink:
			m_curRecord->addXRef( m_text );
			break;
		case Keywords:
			// todo
			break;
		case Source:
			m_curRecord->setAnnotation( c_annSource, m_text );
			break;
		case Organism:
			m_curRecord->setAnnotation( c_annOrganism, m_text );
			break;
		case Lineage:
			m_curRecord->setAnnotation( c_annTaxonomy, m_text );
			break;
		case Authors:
			if( m_curReference )
				m_curReference->setAuthors( m_text );
			break;
		case Title:
			if( m_curReference )
				m_curReference->setTitle( m_text );
			break;
		case Journal:
			if( m_curReference )
				m_curReference->setJournal( m_text );
			break;
		case Pubmed:
			if( m_curReference )
				m_curReference->setDbXRef( QString( "PUBMED:" ) + m_text );
			break;
		case Comment:
			m_curRecord->setComment( m_text );
			break;
		case Features:
			if( m_curFeature )
				endFeature();
			break;
		case Origin:
			m_curRecord->setSequence( new Sequence( m_text.toUpper() ) );
			break;
		default:
			break;
		}

		m_text.clear();
	}

	m_parseMode = mode;

	return true;
}

/**
**/
bool GenBankParser::startReference()
{
	if( m_curReference )
		endReference();

/* Disable reference parsing. Data will be discarded
	m_curReference = new BSeqReference;
*/
	return true;
}

/**
**/
bool GenBankParser::endReference()
{
	if( !m_curRecord || !m_curReference )
		return false;

	m_curRecord->addReference( m_curReference );
	m_curReference = NULL;

	return true;
}

/**
**/
bool GenBankParser::startFeature( const QString &type )
{
	bool	status = true;

	if( m_curFeature )
		status = endFeature();

	m_curFeature = new SeqFeature( type );
	m_featureMode = Location;

	return status;
}

/**
**/
bool GenBankParser::endFeature()
{
	if( !m_curRecord || !m_curFeature )
		return false;

	if( !m_text.isEmpty() ){
		if( !endQualifier() )
			return false;
	}

	m_curRecord->addFeature( m_curFeature );
	m_curFeature = NULL;
	m_featureMode = Invalid;

	return true;
}

/**
**/
bool GenBankParser::startQualifier( char *line )
{
	if( !m_curFeature )
		return false;

	if( m_featureMode == Location ){
		if( !endQualifier() )
			return false;
	}
	else if( m_featureMode == Qualifier ){
		lineWarning( "Unfinished qualifier: " + m_text );
		m_text.clear();
	}

	if( !addText( line, c_featureOffset ) )
		return false;

	int	idx = m_text.indexOf( '=' );

	m_featureMode = Qualifier;

	if( idx < 0 || m_text[idx + 1] != '"' || m_text.endsWith( '"' ) )
		return endQualifier(); // single line qualifier

	return true;
}

/**
**/
bool GenBankParser::endQualifier()
{
	bool	status = true;

	if( m_featureMode == Location )
		status = setLocation();
	else if( m_featureMode == Qualifier )
		status = setQualifier();
	else
		status = false; // not a valid feature mode
	
	m_text.clear();
	m_featureMode = Invalid;

	return status;
}

/**
**/
bool GenBankParser::setLocation()
{
	if( !m_curFeature || m_text.isEmpty() )
		return false;

	SeqRange	location;
		
	if( !parseLocation( m_text, location, SeqLocation::Positive ) )
		return false;

	m_curFeature->setQualifier( "gb_location", m_text );
	m_curFeature->setLocation( location );

	return true;
}

/**
**/
bool GenBankParser::setQualifier()
{
	if( !m_curFeature || m_text.isEmpty() || m_text[0] != '/' )
		return false;

	int		idx = m_text.indexOf( '=' );
	QString	key, value;

	if( idx > 0 ){
		key = m_text.mid( 1, idx - 1 );
		value = m_text.right( m_text.length() - idx - 1 );

		if( key == "translation" )
			value = value.remove( ' ' ); // remove whitespace from wrapped translation qualifiers

		if( value[0] == '"' ){
			int	end = value.length() - 1;
			if( value[end] == '"' )
				value = value.mid( 1, end - 1 );
			else
				lineWarning( "Mismatched double-quotes: " + value );
		}
	}
	else 
		key = m_text.mid( 1 );

	m_curFeature->setQualifier( key, value );

	return true;
}

/**
**/
bool GenBankParser::addText( char *line, int offset )
{
	QByteArray	str;
	
	if( m_parseMode == Origin ){
		str = getAllStripped( line, offset );
	}
	else {
		if( !m_text.isEmpty() && m_featureMode != Location )
			m_text += ' ';

		str = getAll( line, offset );
	}

	m_text += str;

	return true;
}

/**
**/
QByteArray GenBankParser::getAll( char *line, int offset )
{
	line += offset;

	char	*str = line;

	while( *line != '\n' && *line != '\r' && *line != 0 )
		line++;

	*line = 0;

	return QByteArray( str );
}

/**
**/
QByteArray GenBankParser::getAllStripped( char *line, int offset )
{
	QByteArray	res;

	line += offset;

	char	*str = line;

	while( *line != '\n' && *line != '\r' && *line != 0 ){
		if( *line == ' ' || *line == '\t' ){
			if( str != line ){
				*line = 0;
				res += str;
			}
			str = line + 1;
		}
		line++;
	}

	if( str != line ){
		*line = 0;
		res += str;
	}

	return res;
}

/**
**/
QByteArray GenBankParser::getNextToken( char *&line )
{
	char	*str = line;

	while( !isSpace( *line ) && *line != 0 )
		line++;

	if( *line != 0 ){
		*line = 0;
		line++;
		while( isSpace( *line ) )
			line++;
	}

	return QByteArray( str );
}

/**
**/
bool GenBankParser::parseLocus( char *line )
{
	if( !m_curRecord )
		return false;

	int		len;
	QString	str;

	line += c_headerOffset;
	str = getNextToken( line );
	m_curRecord->setName( str );

	len = getNextToken( line ).toInt(); // sequence length
	str = getNextToken( line ); // "bp"
	str = getNextToken( line ); // molecule type
	str = getNextToken( line ); // molecule type pt 2

	str = getNextToken( line ); // division
	m_curRecord->setAnnotation( c_annDataFileDivision, str );

	str = getNextToken( line ); // date
	m_curRecord->setDate( QDate::fromString( str, "dd-MMM-yyyy" ) );

	return true;
}

/**
**/
bool GenBankParser::parseVersion( char *line )
{
	if( !m_curRecord )
		return false;

	QString	ver, gi;
	int		idx;

	line += c_headerOffset;
	ver = getNextToken( line );
	idx = ver.indexOf( '.' );
	if( idx > 0 )
		m_curRecord->setAnnotation( c_annSequenceVersion, ver.right( ver.length() - idx - 1 ) );

	gi = getNextToken( line );
	if( !gi.isEmpty() && gi.startsWith( "GI:" ) )
		m_curRecord->setAnnotation( c_annGenBankId, gi.right( gi.length() - 3 ) );

	return true;
}

/**
**/
bool GenBankParser::parseSource( char *line )
{
	if( !strncmp( "  ORGANISM", line, 10 ) )
		startSection( Organism );
		
	addText( line, c_headerOffset );

	return true;
}

/**
**/
bool GenBankParser::parseOrganism( char *line )
{
	QString	str = getAll( line, c_headerOffset );

	// The species name may wrap lines, so we have to detect the lineage by looking for semicolons
	if( str.contains( ';' ) )
		startSection( Lineage );

	if( !m_text.isEmpty() )
		m_text += ' ';

	m_text += str;

	return true;
}

/**
**/
bool GenBankParser::parseReference( char *line )
{
	if( line[3] != ' ' ) {
		if( !strncmp( "  AUTHORS", line, 9 ) )
			startSection( Authors );
		else if( !strncmp( "  TITLE", line, 7 ) )
			startSection( Title );
		else if( !strncmp( "  JOURNAL", line, 9 ) )
			startSection( Journal );
		else if( !strncmp( "   PUBMED", line, 9 ) )
			startSection( Pubmed );
	}

	return addText( line, c_headerOffset );
}

/**
**/
bool GenBankParser::parseFeatures( char *line )
{
	if( line[5] != ' ' ){
		char	*tok = line + 5;
		startFeature( getNextToken( tok ) );
	}
	else if( m_featureMode != Qualifier && line[c_featureOffset] == '/' )
		return startQualifier( line );

	if( !addText( line, c_featureOffset ) )
		return false;

	if( m_featureMode == Qualifier && m_text.endsWith( '"' ) )
		return endQualifier();

	return true;
}

/**
**/
bool GenBankParser::parseLocation( const QByteArray &str, SeqRange &range, SeqLocation::Strand strand )
{
	if( str[0] == 'j' ){
		if( str.left( 5 ) != "join(" ){
			lineError( "Expected join()." );
			return false;
		}

		int	idx = str.lastIndexOf( ')' );

		if( idx < 0 ){
			lineError( "Missing ')' for join statement." );
			return false;
		}

		if( !parseLocations( str.mid( 5, idx - 5 ), range, strand ) )
			return false;
	}
	else if( str[0] == 'c' ){
		if( str.left( 11 ) != "complement(" ){
			lineError( "Expected complement()." );
			return false;
		}

		int	idx = str.lastIndexOf( ')' );

		if( idx < 0 ){
			lineError( "Missing ')' for complement statement." );
			return false;
		}

		if( strand == SeqLocation::Positive )
			strand = SeqLocation::Negative;
		else if( strand == SeqLocation::Negative )
			strand = SeqLocation::Positive;

		if( !parseLocation( str.mid( 11, idx - 11 ), range, strand ) )
			return false;
	}
	else if( str[0] == 'g' ){
		if( str.left( 6 ) != "group(" ){
			lineError( "Expected group()." );
			return false;
		}

		int	idx = str.lastIndexOf( ')' );

		if( idx < 0 ){
			lineError( "Missing ')' for group statement." );
			return false;
		}
		
		if( !parseLocations( str.mid( 6, idx - 6 ), range, strand ) )
			return false;
	}
	else if( str[0] == 'o' ){
		int	skip = 0;

		if( str.left( 6 ) == "order(" )
			skip = 6;
		else if( str.left( 7 ) == "one-of(" )
			skip = 7;
		else {
			lineError( "Expected order() or one-of()." );
			return false;
		}

		int	idx = str.lastIndexOf( ')' );

		if( idx < 0 ){
			lineError( "Missing ')' for statement." );
			return false;
		}
		
		if( !parseLocations( str.mid( skip, idx - skip ), range, strand ) )
			return false;
	}
	else {
		int	start = 0, end = 0;
		int	idx = str.indexOf( ':' );

		if( idx > 0 ){
			QString	ref = str.left( idx );

			if( !parseStartEnd( str.mid( idx + 1 ), start, end ) )
				return false;

			SeqLocation	loc( start, end, strand );

			loc.setReference( ref );
			range.addLocation( loc );
		}
		else {
			if( !parseStartEnd( str, start, end ) )
				return false;

			range.addLocation( SeqLocation( start, end, strand ) );
		}
	}

	return true;
}

/**
**/
bool GenBankParser::parseLocations( const QByteArray &str, SeqRange &range, SeqLocation::Strand strand )
{
	int	idx = 0;
	int	cur = 0;
	int	end = str.length();

	while( cur < end ){
		if( str[cur] == ',' ){
			if( !parseLocation( str.mid( idx, cur - idx ), range, strand ) )
				return false;
			cur++;
			idx = cur;
		}
		else if( str[cur] == '(' ){
			int	parenLvl = 1;

			cur++;
			while( cur < end && parenLvl > 0 ){
				if( str[cur] == '(' )
					parenLvl++;
				else if( str[cur] == ')' )
					parenLvl--;
				cur++;
			}
		}
		else
			cur++;
	}

	if( idx < end ){
		if( !parseLocation( str.mid( idx ), range, strand ) )
			return false;
	}

	return true;
}

/**
**/
bool GenBankParser::parseStartEnd( const QByteArray &str, int &start, int &end )
{
	const char	*data = str.data();

	while( *data != 0 && *data == ' ' )
		data++;

	if( *data == '<' )
		data++;
	else if( *data == '>' )
		data++;

	if( *data == '(' ){
		int	temp;
		int	idx = str.indexOf( ')' );

		if( idx < 0 ){
			lineError( "Unmatched parentheses in location" );
			return false;
		}

		data++;
		if( !parseStartEnd( QByteArray( data ), start, temp ) )
			return false;

		data = str.data() + idx + 1;
	}
	else {
		while( *data != 0 && *data >= '0' && *data <= '9' ){
			start *= 10;
			start += *data - '0';
			data++;
		}
	}

	if( *data == 0 )
		end = start;
	else {
		if( *data == '.' ){
			data++;

			if( *data == '.' )
				data++;
		}
		else if( *data == '^' )
			data++;
		else if( *data == '>' )
			data++;
		else
			lineError( "Expected '.', '..', '>' or '^'" );
	}

	if( *data == '(' ){
		int	temp;

		data++;
		if( !parseStartEnd( QByteArray( data ), temp, end ) )
			return false;
	}
	else {
		while( *data != 0 && *data >= '0' && *data <= '9' ){
			end *= 10;
			end += *data - '0';
			data++;
		}

		if( *data != 0 && *data != '>' && *data != ')' )
			lineError( "Expected end of sequence." );
	}

	return true;
}
